"""
邮件拉取优化模块
功能：批量FETCH、增量拉取、批量数据库插入
"""

import email
import json
from datetime import datetime, timedelta
from email.header import decode_header
import logging

logger = logging.getLogger('email_system')


def safe_decode_header(header_value):
    """安全解码邮件头部"""
    if not header_value:
        return ''
    
    try:
        decoded_parts = decode_header(header_value)
        result = []
        
        for content, encoding in decoded_parts:
            if isinstance(content, bytes):
                if encoding:
                    result.append(content.decode(encoding, errors='ignore'))
                else:
                    result.append(content.decode('utf-8', errors='ignore'))
            else:
                result.append(str(content))
        
        return ''.join(result)
    except Exception as e:
        logger.error(f"解码头部失败: {e}")
        return str(header_value)


def batch_fetch_headers(mail, email_ids, account, folder_name, batch_size=30):
    """
    批量拉取邮件头部信息
    
    Args:
        mail: IMAP连接
        email_ids: 邮件ID列表
        account: 邮箱账号对象
        folder_name: 文件夹名称
        batch_size: 每批处理数量
    
    Returns:
        邮件头部数据列表
    """
    from models import EmailMessage
    
    headers_list = []
    total = len(email_ids)
    
    logger.info(f"[批量拉取] 开始拉取 {total} 封邮件，批次大小: {batch_size}")
    
    for i in range(0, total, batch_size):
        batch_ids = email_ids[i:i + batch_size]
        batch_num = i // batch_size + 1
        
        try:
            # 🚀 关键优化：批量FETCH
            id_list = b','.join(batch_ids)
            status, msg_data_list = mail.fetch(id_list, '(UID RFC822.SIZE RFC822.HEADER)')
            
            if status != 'OK' or not msg_data_list:
                logger.warning(f"[批量拉取] 批次 {batch_num} 失败")
                continue
            
            # 解析批量返回的数据
            parsed_count = 0
            for j in range(0, len(msg_data_list), 2):
                if j + 1 >= len(msg_data_list):
                    break
                
                try:
                    item = msg_data_list[j]
                    
                    if not isinstance(item, tuple) or len(item) < 2:
                        continue
                    
                    header_str = item[0].decode('utf-8', errors='ignore')
                    header_content = item[1]
                    
                    # 解析UID
                    uid = None
                    if 'UID' in header_str:
                        try:
                            uid_part = header_str.split('UID')[1].split()[0].rstrip(')')
                            uid = int(uid_part)
                        except (ValueError, IndexError):
                            continue
                    
                    # 解析大小
                    size = None
                    if 'RFC822.SIZE' in header_str:
                        try:
                            size_part = header_str.split('RFC822.SIZE')[1].split()[0]
                            size = int(size_part)
                        except (ValueError, IndexError):
                            pass
                    
                    if not isinstance(header_content, bytes):
                        continue
                    
                    # 解析邮件头部
                    email_message = email.message_from_bytes(header_content)
                    
                    # 提取Message-ID
                    message_id = email_message.get('Message-ID', f"local-{uid or i+j}")
                    
                    # 检查是否已存在（避免重复）
                    existing = EmailMessage.query.filter_by(
                        user_id=account.user_id,
                        message_id=message_id
                    ).first()
                    
                    if existing:
                        continue
                    
                    # 解析邮件头部字段
                    subject = safe_decode_header(email_message.get('Subject', '(无主题)'))
                    sender = safe_decode_header(email_message.get('From', ''))
                    recipients = safe_decode_header(email_message.get('To', ''))
                    cc = safe_decode_header(email_message.get('Cc', ''))
                    
                    # 解析接收时间
                    received_at = datetime.utcnow()
                    try:
                        date_str = email_message.get('Date')
                        if date_str:
                            from email.utils import parsedate_to_datetime
                            parsed_date = parsedate_to_datetime(date_str)
                            received_at = parsed_date.astimezone().replace(tzinfo=None)
                    except:
                        pass
                    
                    # 创建邮件记录
                    message_data = {
                        'user_id': account.user_id,
                        'email_account_id': account.id,
                        'message_id': message_id,
                        'imap_uid': uid,
                        'subject': subject,
                        'sender': sender,
                        'recipients': json.dumps([recipients]) if recipients else '[]',
                        'cc': json.dumps([cc]) if cc else '[]',
                        'body_text': '',
                        'body_html': '',
                        'attachments': '[]',
                        'is_full_content': False,
                        'content_size': size,
                        'folder': folder_name.lower(),
                        'received_at': received_at
                    }
                    
                    headers_list.append(message_data)
                    parsed_count += 1
                
                except Exception as e:
                    logger.error(f"[批量拉取] 解析单封邮件失败: {e}")
                    continue
            
            logger.info(f"[批量拉取] 批次 {batch_num}/{(total + batch_size - 1) // batch_size} 完成，解析 {parsed_count} 封")
        
        except Exception as e:
            logger.error(f"[批量拉取] 批次 {batch_num} 异常: {e}")
            continue
    
    logger.info(f"[批量拉取] 总计成功拉取 {len(headers_list)} 封邮件头部")
    return headers_list


def get_fetch_strategy(account, folder_name):
    """
    获取拉取策略（增量或首次）
    
    Returns:
        (is_incremental, search_criteria)
    """
    from models import EmailMessage
    from config import Config
    
    # 检查是否有历史记录
    latest_message = EmailMessage.query.filter_by(
        email_account_id=account.id,
        folder=folder_name.lower()
    ).order_by(EmailMessage.received_at.desc()).first()
    
    if latest_message and latest_message.received_at:
        # 增量拉取：拉取最后一封邮件之后的
        since_date = latest_message.received_at - timedelta(hours=1)  # 留1小时缓冲
        date_str = since_date.strftime('%d-%b-%Y')
        logger.info(f"[拉取策略] 增量拉取 - 文件夹: {folder_name}, 起始日期: {date_str}")
        return True, f'SINCE {date_str}'
    else:
        # 首次拉取：只拉取最近7天
        days_ago = datetime.now() - timedelta(days=Config.EMAIL_INITIAL_FETCH_DAYS)
        date_str = days_ago.strftime('%d-%b-%Y')
        logger.info(f"[拉取策略] 首次拉取 - 文件夹: {folder_name}, 起始日期: {date_str}")
        return False, f'SINCE {date_str}'


def fetch_folder_emails(mail, account, folder_info):
    """
    拉取单个文件夹的邮件
    
    Args:
        mail: IMAP连接
        account: 邮箱账号
        folder_info: (imap_name, display_name) 元组
    
    Returns:
        拉取的邮件数量
    """
    from models import db, EmailMessage
    from config import Config
    
    imap_name, display_name = folder_info
    
    try:
        # 选择文件夹
        status, count_data = mail.select(imap_name)
        if status != 'OK':
            logger.warning(f"[拉取] 选择文件夹失败 - {display_name}")
            return 0
        
        # 获取拉取策略
        is_incremental, search_criteria = get_fetch_strategy(account, display_name)
        
        # 搜索邮件
        status, messages = mail.search(None, search_criteria)
        
        if status != 'OK' or not messages[0]:
            # 如果按日期搜索失败，尝试获取最新的邮件
            if not is_incremental:
                status, messages = mail.search(None, 'ALL')
                if status == 'OK' and messages[0]:
                    email_ids = messages[0].split()
                    # 只取最新的50封
                    email_ids = email_ids[-Config.EMAIL_INITIAL_FETCH_COUNT:]
                else:
                    logger.info(f"[拉取] 文件夹 {display_name} 没有邮件")
                    return 0
            else:
                logger.info(f"[拉取] 文件夹 {display_name} 没有新邮件")
                return 0
        else:
            email_ids = messages[0].split()
            
            # 首次拉取限制数量
            if not is_incremental and len(email_ids) > Config.EMAIL_INITIAL_FETCH_COUNT:
                email_ids = email_ids[-Config.EMAIL_INITIAL_FETCH_COUNT:]
        
        if not email_ids:
            logger.info(f"[拉取] 文件夹 {display_name} 没有邮件")
            return 0
        
        logger.info(f"[拉取] 文件夹 {display_name} 将拉取 {len(email_ids)} 封邮件")
        
        # 批量拉取邮件头部
        headers_list = batch_fetch_headers(mail, email_ids, account, display_name)
        
        if not headers_list:
            logger.info(f"[拉取] 文件夹 {display_name} 没有新邮件")
            return 0
        
        # 🚀 批量插入数据库
        db.session.bulk_insert_mappings(EmailMessage, headers_list)
        db.session.commit()
        
        logger.info(f"[拉取] 文件夹 {display_name} 完成 - 新增 {len(headers_list)} 封邮件")
        return len(headers_list)
    
    except Exception as e:
        logger.error(f"[拉取] 文件夹 {display_name} 失败: {e}")
        db.session.rollback()
        return 0
